///@file PCbasis.cpp
///@brief Implementation of a polynomial chaos basis
///@author Arnaud Duval
///@version 0.0
///@date 2010/05/10

#include "PCbasis.h"
#include "ProbUtils.h"
#include "IndexSet.h"
#include "StringUtils.h"

///@brief Default constructor
PCbasis::PCbasis()
{
	M = 0;
}

///@brief Constructor with arguments
///@param DM number of variables
///@param types List of variables types
///@param p maximum degree of polynomials
PCbasis::PCbasis(unsigned int DM, std::string* types, unsigned int p)
{
	M = DM;
	P = Binomial(M+p, p);
	poly = new Polynomial*[P];
	for(unsigned int i = 0 ; i < P ; i++)
		poly[i] = new Polynomial[M];
	
	idx = IndexSet(M,p);

	for(unsigned int i = 0 ; i < P ; i++)
	{
		for(unsigned int j = 0 ; j < M ; j++)
		{
			if(types[j] == "gauss")
			{
				poly[i][j] = Polynomial(idx[i][j+1], 'H');
			}
			else if(types[j] == "uniform")
			{
				poly[i][j] = Polynomial(idx[i][j+1], 'L');
			}
		}
	}
}

///@brief Copy constructor
///@param p PCbasis object to copy
PCbasis::PCbasis(const PCbasis& p)
{
	M = p.M;
	P = p.P;
	idx = p.idx;
	poly = new Polynomial*[P];
	for(unsigned int i = 0 ; i < P ; i++)
		poly[i] = new Polynomial[M];
	for(unsigned int i = 0 ; i < P ; i++)
	{
		for(unsigned int j = 0 ; j < M ; j++)
		{
			poly[i][j] = p.poly[i][j];
		}
	}
}

///@brief Copy constructor from temporary object
///@param p pointer to PCbasis object to copy
PCbasis::PCbasis(PCbasis* p)
{
	M = (*p).M;
	P = (*p).P;
	idx = (*p).idx;
	poly = new Polynomial*[P];
	for(unsigned int i = 0 ; i < P ; i++)
		poly[i] = new Polynomial[M];
	for(unsigned int i = 0 ; i < P ; i++)
	{
		for(unsigned int j = 0 ; j < M ; j++)
		{
			poly[i][j] = (*p).poly[i][j];
		}
	}
}

///@brief Constructor with truncation from an existing basis
///@param b Existing basis
///@param id index of the removed MultiIndex
PCbasis::PCbasis(const PCbasis& b, unsigned int id)
{
	M = b.M;
	P = b.P - 1;
	idx = b.idx;
	assert(id >= 0);
	assert(id < b.P);
	idx.Remove(id);
	poly = new Polynomial*[P];
	for(unsigned int i = 0 ; i < P ; i++)
		poly[i] = new Polynomial[M];
	for(unsigned int i = 0 ; i < P ; i++)
	{
		for(unsigned int j = 0 ; j < M ; j++)
		{
			if(i < id)
				poly[i][j] = b.poly[i][j];
			else
				poly[i][j] = b.poly[i+1][j];
		}
	}
}

///@brief Constructor with a given IndexSet
///@param s Given IndexSet
///@param types List of variables types
PCbasis::PCbasis(IndexSet s, std::string* types)
{
	M = s[0].Cardinality();
	P = s.Cardinality();
	poly = new Polynomial*[P];
	for(unsigned int i = 0 ; i < P ; i++)
		poly[i] = new Polynomial[M];
	
	idx = s;

	for(unsigned int i = 0 ; i < P ; i++)
	{
		for(unsigned int j = 0 ; j < M ; j++)
		{
			if(types[j] == "gauss")
			{
				poly[i][j] = Polynomial(idx[i][j+1], 'H');
			}
			else if(types[j] == "uniform")
			{
				poly[i][j] = Polynomial(idx[i][j+1], 'L');
			}
		}
	}
}

///@brief Destructor
PCbasis::~PCbasis()
{
	for(unsigned int i = 0 ; i < P ; i++)
		delete[] poly[i];
	delete[] poly;
}

///@brief Overload of = operator
///@param p PCbasis object to copy
PCbasis& PCbasis::operator=(const PCbasis& p)
{
	M = p.M;
	P = p.P;
	idx = p.idx;
	poly = new Polynomial*[P];
	for(unsigned int i = 0 ; i < P ; i++)
		poly[i] = new Polynomial[M];
	for(unsigned int i = 0 ; i < P ; i++)
	{
		for(unsigned int j = 0 ; j < M ; j++)
		{
			poly[i][j] = p.poly[i][j];
		}
	}
	return (*this);
}
///@brief Overload of = operator for temporary object
///@param p Pointer to PCbasis object to copy
PCbasis& PCbasis::operator=(PCbasis* p)
{
	M = (*p).M;
	P = (*p).P;
	idx = (*p).idx;
	poly = new Polynomial*[P];
	for(unsigned int i = 0 ; i < P ; i++)
		poly[i] = new Polynomial[M];
	for(unsigned int i = 0 ; i < P ; i++)
	{
		for(unsigned int j = 0 ; j < M ; j++)
		{
			poly[i][j] = (*p).poly[i][j];
		}
	}
	return (*this);
}

///@brief Return PCbasis cardinality
unsigned int PCbasis::Cardinality()
{
	return P;
}

///@brief Compute value with a given variables vector
Vector PCbasis::Evaluate(Vector x)
{
	Vector res(P);
	for(unsigned int i = 1 ; i <= P ; i++)
	{
		res[i] = 1.;
		for(unsigned int j = 0 ; j < M ; j++)
			res[i] *= poly[i-1][j].Evaluate(x[j+1]);
	}
	return res;
}

///@brief Return IndexSet used for PC basis
IndexSet PCbasis::A()
{
	return idx;
}

///@brief Return number of variables
unsigned int PCbasis::_M()
{
	return M;
}

///@brief Overload of << operator
///@param s out stream
///@param b PCbasis to put in the out stream
std::ostream& operator<<(std::ostream &s, const PCbasis& b)
{
	for(unsigned int i = 0 ; i < b.P ; i++)
	{
		s << i+1 << " : ";
		std::string e;
		for(unsigned int j = 0 ; j < b.M ; j++)
		{
			std::string var = "x_"+IntToString(j+1, 1);
			e += "("+b.poly[i][j].Expression(var)+")";
			if(j != b.M-1)
				e += " * ";
		}
		s << e << "\n";
	}
	return s;
}
